import cv2
import pytesseract
from image_processing import enhance_image, extract_paper, rotate_image_180_flip

def clean_and_join_lines(text):
    """
    清理并拼接提取的文字中的换行符。
    :param text: 提取出的文字内容
    :return: 拼接后更易读的文字内容
    """
    lines = text.split('\n')  # 按行分割
    cleaned_lines = []
    current_line = ""
    for line in lines:
        stripped_line = line.strip()  # 去除首尾空格
        if stripped_line:
            if current_line:
                current_line += stripped_line  # 拼接当前行
            else:
                current_line = stripped_line  # 开始新的一行
        else:
            if current_line:
                cleaned_lines.append(current_line)  # 保存当前行
                current_line = ""
    if current_line:
        cleaned_lines.append(current_line)  # 处理最后一行
    return "\n".join(cleaned_lines)  # 返回拼接后的文本

def identify_address(enhanced_image, extract_image, address_extract_flag):
    """
    识别地址区域中的文本。
    :param enhanced_image: 增强后的图像
    :param extract_image: 提取的图像
    :param address_extract_flag: 用于识别地址的关键字
    :return: 识别出的地址文本
    """
    data = pytesseract.image_to_data(enhanced_image, lang='chi_sim', output_type=pytesseract.Output.DICT)  # 使用Tesseract识别文字及其位置
    for i in range(len(data['text'])):
        print(data['text'][i])
        if address_extract_flag in data['text'][i]:  # 查找包含地址关键字的文本
            (x, y, w, h) = (data['left'][i], data['top'][i], data['width'][i], data['height'][i])  # 获取边界框
            height, width = extract_image.shape[:2]
            # 根据图片的宽高比选择抠图区域
            address_area = extract_image[y - 10:y + h + 200, x - 100:x + w + 10] if height < width else extract_image[y - 15:y + h + 100, x - 10:x + w + 300]
            cv2.imwrite("extracted_address_area.png", address_area)  # 保存抠出的地址区域
            address_text = pytesseract.image_to_string(address_area, lang='chi_sim')  # 识别地址区域的文字
            cleaned_text = clean_and_join_lines(address_text)  # 清理并拼接识别出的文本
            return cleaned_text
    return None

def extract_text_line_image(enhanced_image, extract_image, address_extract_flag):
    """
    识别地址区域中的文本，并提取包含关键字的那一行对应的图像。
    :param enhanced_image: 增强后的图像
    :param extract_image: 提取的图像
    :param address_extract_flag: 用于识别地址的关键字
    :return: 包含关键字那一行的图像
    """
    data = pytesseract.image_to_data(enhanced_image, lang='chi_sim', output_type=pytesseract.Output.DICT)
    height, width = extract_image.shape[:2]

    for i in range(len(data['text'])):
        if address_extract_flag in data['text'][i]:
            # 获取关键字所在的行索引
            baseline_y = data['top'][i]

            # 计算该行起始和结束的x坐标
            line_start_x = width
            line_end_x = 0
            padding = 10  # 裁剪时上下左右增加一点边缘

            # 向左寻找行的起点
            for j in range(i, -1, -1):
                if ' '.join(data['text'][j:j + 1]).strip():
                    line_start_x = min(line_start_x, data['left'][j])
                else:
                    break

            # 向右寻找行的终点
            for k in range(i, len(data['text'])):
                if ' '.join(data['text'][k:k + 1]).strip():
                    line_end_x = max(line_end_x, data['left'][k] + data['width'][k])
                else:
                    break

            # 根据找到的坐标裁剪图像
            # 调整y方向上的边界以减少偏移
            y_start = max(baseline_y - padding, 0)  # 减少向上偏移
            y_end = min(baseline_y + data['height'][i] + padding, height)
            extracted_line_image = extract_image[y_start:y_end, line_start_x:line_end_x]

            cv2.imwrite("extracted_text_line.png", extracted_line_image)  # 保存提取的行图像
            return extracted_line_image  # 返回提取的图像数据
    return None


def extract_text_lines_image(enhanced_image, extract_image, address_extract_flag):
    """
    识别地址区域中的文本，并提取包含关键字的所有行对应的图像。
    :param enhanced_image: 增强后的图像
    :param extract_image: 提取的图像
    :param address_extract_flag: 用于识别地址的关键字
    :return: 包含关键字所有行的图像列表
    """
    data = pytesseract.image_to_data(enhanced_image, lang='chi_sim', output_type=pytesseract.Output.DICT)
    height, width = extract_image.shape[:2]

    # 初始化变量
    extracted_line_images = []
    current_line = {"start_x": width, "end_x": 0, "top_y": height, "bottom_y": 0}
    found_keyword = False

    for i in range(len(data['text'])):
        if address_extract_flag in data['text'][i].strip():
            found_keyword = True
            x = data['left'][i]
            y = data['top'][i]
            w = data['width'][i]
            h = data['height'][i]

            # 更新当前行的边界
            current_line["start_x"] = min(current_line["start_x"], x)
            current_line["end_x"] = max(current_line["end_x"], x + w)
            current_line["top_y"] = min(current_line["top_y"], y)
            current_line["bottom_y"] = max(current_line["bottom_y"], y + h)

            # 向前合并直到空白行
            j = i - 1
            while j >= 0 and ' '.join(data['text'][j:j + 1]).strip():
                x_left = data['left'][j]
                w_left = data['width'][j]
                current_line["start_x"] = min(current_line["start_x"], x_left)
                current_line["top_y"] = min(current_line["top_y"], data['top'][j])
                current_line["bottom_y"] = max(current_line["bottom_y"], data['top'][j] + data['height'][j])
                j -= 1

            # 向后合并直到空白行
            k = i + 1
            while k < len(data['text']) and ' '.join(data['text'][k:k + 1]).strip():
                x_right = data['left'][k]
                w_right = data['width'][k]
                current_line["end_x"] = max(current_line["end_x"], x_right + w_right)
                current_line["top_y"] = min(current_line["top_y"], data['top'][k])
                current_line["bottom_y"] = max(current_line["bottom_y"], data['top'][k] + data['height'][k])
                k += 1

        elif found_keyword:
            # 添加边缘填充
            padding = 5
            start_x = max(current_line["start_x"] - padding, 0)
            end_x = min(current_line["end_x"] + padding, width)
            top_y = max(current_line["top_y"] - padding, 0)
            bottom_y = min(current_line["bottom_y"] + padding, height)

            # 裁剪并保存图像
            extracted_line_image = extract_image[top_y:bottom_y, start_x:end_x]
            extracted_line_images.append(extracted_line_image)
            cv2.imwrite(f"extracted_text_line_{len(extracted_line_images)}.png", extracted_line_image)

            # 重置变量
            current_line = {"start_x": width, "end_x": 0, "top_y": height, "bottom_y": 0}
            found_keyword = False

    # 如果最后一行有关键字且未被保存，则保存它
    if found_keyword:
        padding = 5
        start_x = max(current_line["start_x"] - padding, 0)
        end_x = min(current_line["end_x"] + padding, width)
        top_y = max(current_line["top_y"] - padding, 0)
        bottom_y = min(current_line["bottom_y"] + padding, height)

        extracted_line_image = extract_image[top_y:bottom_y, start_x:end_x]
        extracted_line_images.append(extracted_line_image)
        cv2.imwrite(f"extracted_text_line_{len(extracted_line_images)}.png", extracted_line_image)

    return extracted_line_images

def  extract_text_lines_image1(enhanced_image, extract_image, address_extract_flag):
    """
    识别地址区域中的文本，并提取包含关键字的所有行对应的图像。
    :param enhanced_image: 增强后的图像
    :param extract_image: 提取的图像
    :param address_extract_flag: 用于识别地址的关键字
    :return: 包含关键字所有行的图像列表
    """
    data = pytesseract.image_to_data(enhanced_image, lang='chi_sim', output_type=pytesseract.Output.DICT)
    height, width = extract_image.shape[:2]

    extracted_line_images = []
    current_line = None

    # 遍历所有的文本块
    for i in range(len(data['text'])):
        text = data['text'][i].strip()
        if not text:
            continue  # 跳过空文本

        x = data['left'][i]
        y = data['top'][i]
        w = data['width'][i]
        h = data['height'][i]

        if address_extract_flag in text:
            # found_keyword = True
            if current_line is None:
                current_line = {"start_x": x, "end_x": x + w, "top_y": y, "bottom_y": y + h}
            else:
                current_line["start_x"] = min(current_line["start_x"], x)
                current_line["end_x"] = max(current_line["end_x"], x + w)
                current_line["top_y"] = min(current_line["top_y"], y)
                current_line["bottom_y"] = max(current_line["bottom_y"], y + h)

            # 向前合并直到空白行或不同的块
            j = i - 1
            while j >= 0 and data['text'][j].strip() and data['block_num'][j] == data['block_num'][i]:
                x_left = data['left'][j]
                w_left = data['width'][j]
                current_line["start_x"] = min(current_line["start_x"], x_left)
                current_line["top_y"] = min(current_line["top_y"], data['top'][j])
                current_line["bottom_y"] = max(current_line["bottom_y"], data['top'][j] + data['height'][j])
                j -= 1

            # 向后合并直到空白行或不同的块
            k = i + 1
            while k < len(data['text']) and data['text'][k].strip() and data['block_num'][k] == data['block_num'][i]:
                x_right = data['left'][k]
                w_right = data['width'][k]
                current_line["end_x"] = max(current_line["end_x"], x_right + w_right)
                current_line["top_y"] = min(current_line["top_y"], data['top'][k])
                current_line["bottom_y"] = max(current_line["bottom_y"], data['top'][k] + data['height'][k])
                k += 1

            padding = 0
            start_x = max(current_line["start_x"] - padding, 0)
            end_x = min(current_line["end_x"] + padding, width)
            top_y = max(current_line["top_y"] - padding, 0)
            bottom_y = min(current_line["bottom_y"] + padding, height)

            extracted_line_image = extract_image[top_y:bottom_y, start_x:end_x]
            extracted_line_images.append(extracted_line_image)
            cv2.imwrite(f"extracted_text_line_{len(extracted_line_images)}.png", extracted_line_image)

            # # 重置变量
            current_line = None
    return extracted_line_images

def extract_text_area_image(enhanced_image, extract_image, address_extract_flag):
    """
    识别地址区域内所有包含关键字的文本行，并合并这些行对应的图像。
    :param enhanced_image: 增强后的图像
    :param extract_image: 提取的图像
    :param address_extract_flag: 用于识别地址的关键字
    :return: 包含所有关键字行的图像区域
    """
    data = pytesseract.image_to_data(enhanced_image, lang='chi_sim', output_type=pytesseract.Output.DICT)
    height, width = extract_image.shape[:2]

    # 存储所有包含关键字的行的边界信息
    keyword_lines = []
    padding = 10  # 裁剪时上下左右增加一点边缘

    for i in range(len(data['text'])):
        if address_extract_flag in data['text'][i].strip():
            baseline_y = data['top'][i]
            line_height = data['height'][i]

            # 计算当前行的left和right
            line_start_x = data['left'][i]
            line_end_x = data['left'][i] + data['width'][i]

            # 向左寻找行的起点
            for j in range(i, -1, -1):
                if ' '.join(data['text'][j:j + 1]).strip():
                    line_start_x = min(line_start_x, data['left'][j])
                else:
                    break

            # 向右寻找行的终点
            for k in range(i, len(data['text'])):
                if ' '.join(data['text'][k:k + 1]).strip():
                    line_end_x = max(line_end_x, data['left'][k] + data['width'][k])
                else:
                    break

            # 添加或更新当前行的信息
            found_line = False
            for line_info in keyword_lines:
                if (baseline_y >= line_info['top'] and baseline_y <= line_info['bottom']) or \
                        ((baseline_y + line_height) >= line_info['top'] and (baseline_y + line_height) <= line_info[
                            'bottom']):
                    line_info['top'] = min(line_info['top'], max(baseline_y - padding, 0))
                    line_info['bottom'] = max(line_info['bottom'], min(baseline_y + line_height + padding, height))
                    line_info['left'] = min(line_info['left'], line_start_x)
                    line_info['right'] = max(line_info['right'], line_end_x)
                    found_line = True
                    break

            if not found_line:
                keyword_lines.append({
                    'top': max(baseline_y - padding, 0),
                    'bottom': min(baseline_y + line_height + padding, height),
                    'left': line_start_x,
                    'right': line_end_x
                })

    # 如果找到了包含关键字的行，则计算整个区域的边界并裁剪
    if keyword_lines:
        top = min(item['top'] for item in keyword_lines)
        bottom = max(item['bottom'] for item in keyword_lines)
        left = min(item['left'] for item in keyword_lines)
        right = max(item['right'] for item in keyword_lines)

        extracted_area_image = extract_image[top:bottom, left:right]
        cv2.imwrite("extracted_text_area.png", extracted_area_image)  # 保存提取的区域图像
        return extracted_area_image
    else:
        return None

def extract_address_area(image_path, extracted_paper_path, address_extract_flag):
    """
    提取地址区域。
    :param image_path: 输入图像的路径
    :param extracted_paper_path: 提取后的纸张图像保存路径
    :param address_extract_flag: 用于识别地址的关键字
    :return: 识别出的地址文本
    """
    extract_paper(image_path, extracted_paper_path)  # 提取纸张
    enhanced_image = enhance_image(extracted_paper_path)  # 增强纸张
    extract_image = cv2.imread(extracted_paper_path)  # 读取增强后的纸张
    address = extract_text_lines_image1(enhanced_image, extract_image, address_extract_flag)  # 识别地址
    if address:
        return address
    rotate_image_180_flip(extract_image, extracted_paper_path)  # 旋转180度
    enhanced_image = enhance_image(extracted_paper_path)  # 再次增强纸张
    image = cv2.imread(extracted_paper_path)  # 读取旋转后的纸张
    address = extract_text_lines_image1(enhanced_image, image, address_extract_flag)  # 识别地址
    return address